// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/dwio/common/seekable_input_stream.h>
#include <pollux/dwio/common/compression/compression.h>

namespace kumo::pollux::dwio::common::compression {
    class PagedInputStream : public dwio::common::SeekableInputStream {
    public:
        PagedInputStream(
            std::unique_ptr<SeekableInputStream> inStream,
            memory::MemoryPool &memPool,
            std::unique_ptr<Decompressor> decompressor,
            const dwio::common::encryption::Decrypter *decrypter,
            const std::string &streamDebugInfo,
            bool useRawDecompression = false,
            size_t compressedLength = 0)
            : input_(std::move(inStream)),
              pool_(memPool),
              inputBuffer_(pool_),
              decompressor_{std::move(decompressor)},
              decrypter_{decrypter},
              streamDebugInfo_{streamDebugInfo} {
            DWIO_ENSURE(
                decompressor_ || decrypter_,
                "one of decompressor or decryptor is required");
            DWIO_ENSURE(
                !useRawDecompression || compressedLength > 0,
                "For raw decompression, compressedLength should be greater than zero");

            if (useRawDecompression) {
                state_ = State::START;
                remainingLength_ = compressedLength;
            }
        }

        bool Next(const void **data, int32_t *size) override;

        void BackUp(int32_t count) override;

        // NOTE: This always returns true.
        bool SkipInt64(int64_t count) override;

        google::protobuf::int64 ByteCount() const override {
            return bytesReturned_ + pendingSkip_;
        }

        void seekToPosition(dwio::common::PositionProvider &position) override;

        std::string getName() const override {
            return melon::to<std::string>(
                "PagedInputStream StreamInfo (",
                streamDebugInfo_,
                ") input stream (",
                input_->getName(),
                ") State (",
                state_,
                ") remaining length (",
                remainingLength_,
                ")");
        }

        size_t positionSize() const override {
            // not compressed, so need 2 positions (compressed position + uncompressed
            // position)
            return 2;
        }

    protected:
        // Special constructor used by ZlibDecompressionStream
        PagedInputStream(
            std::unique_ptr<SeekableInputStream> inStream,
            memory::MemoryPool &memPool,
            const std::string &streamDebugInfo,
            bool useRawDecompression = false,
            size_t compressedLength = 0)
            : input_(std::move(inStream)),
              pool_(memPool),
              inputBuffer_(pool_),
              decompressor_{nullptr},
              decrypter_{nullptr},
              streamDebugInfo_{streamDebugInfo} {
            DWIO_ENSURE(
                !useRawDecompression || compressedLength > 0,
                "For raw decompression, compressedLength should be greater than zero");

            if (useRawDecompression) {
                state_ = State::START;
                remainingLength_ = compressedLength;
            }
        }

        void prepareOutputBuffer(uint64_t uncompressedLength);

        void readBuffer(bool failOnEof);

        uint32_t readByte(bool failOnEof);

        void readHeader();

        void clearDecompressionState();

        enum class State { HEADER, START, ORIGINAL, END };

        // make sure input is contiguous for decompression/decryption
        const char *ensureInput(size_t availableInputBytes);

        virtual bool readOrSkip(const void **data, int32_t *size);

        // input stream where to read compressed/encrypted data
        std::unique_ptr<SeekableInputStream> input_;
        memory::MemoryPool &pool_;

        // Offset in input_ of the last header read
        uint64_t lastHeaderOffset_{0};

        // The value of bytesReturned_ at the time when last header was read
        uint64_t bytesReturnedAtLastHeaderOffset_{0};

        // buffer to hold an entire compressed/encrypted block allowing
        // decompression/decryption algorithm to work on contiguous block
        dwio::common::DataBuffer<char> inputBuffer_;

        // uncompressed output
        std::unique_ptr<dwio::common::DataBuffer<char> > outputBuffer_{nullptr};

        // unencrypted output
        std::unique_ptr<melon::IOBuf> decryptionBuffer_{nullptr};

        // the current state
        State state_{State::HEADER};

        // the start of the current output buffer
        const char *outputBufferPtr_{nullptr};

        // the size of the current output buffer
        size_t outputBufferLength_{0};

        // the size of the current chunk (in its compressed/encrypted form)
        size_t remainingLength_{0};

        // The first byte in the range from last call to 'input_->Next()'.
        const char *inputBufferStart_{nullptr};

        // The first byte to return in Next. Not the same as inputBufferStart_ if
        // there has been a BackUp().
        const char *inputBufferPtr_{nullptr};

        // The first byte after the last range returned by 'input_->Next()'.
        const char *inputBufferPtrEnd_{nullptr};

        // Bytes returned or skipped by this stream, not including pendingSkip_.
        uint64_t bytesReturned_{0};

        // Size returned by the previous call to Next().
        int32_t lastWindowSize_{0};

        // decompressor
        std::unique_ptr<Decompressor> decompressor_;

        // decrypter
        const dwio::common::encryption::Decrypter *decrypter_;

        int64_t pendingSkip_{0};

    private:
        bool skipAllPending();

        // Stream Debug Info
        const std::string streamDebugInfo_;
    };
} // namespace kumo::pollux::dwio::common::compression
